import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import os
import asyncio
import threading
from docx import Document
from PyPDF2 import PdfReader
from templates import (
    get_template_names, get_template, combine_templates,
    ANALYST_ROLES, combine_roles
)
from datetime import datetime
from api import TongYiAPI
import json

class BatchAnalyzerGUI:
    def __init__(self, root):
        self.root = root
        self.api = TongYiAPI()
        
        # 加载保存的设置
        self.load_settings()
        
        # 初始化变量
        self.results = {}
        self.is_running = False
        self.file_paths = []
        
        # 初始化角色选择
        self.main_role_var = tk.StringVar(value="首席策略分析师")
        self.selected_assistant_roles = []  # 存储选中的辅助角色
        
        # 设置窗口大小和位置
        window_width = 1200
        window_height = 700
        screen_width = root.winfo_screenwidth()
        screen_height = root.winfo_screenheight()
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2
        root.geometry(f"{window_width}x{window_height}+{x}+{y}")
        root.minsize(1000, 600)
        
        self.setup_ui()
    
    def setup_ui(self):
        # 创建左右分栏
        container = ttk.Frame(self.root)
        container.grid(row=0, column=0, sticky='nsew', padx=5, pady=5)
        
        # 配置根窗口的网格权重
        self.root.grid_rowconfigure(0, weight=1)
        self.root.grid_columnconfigure(0, weight=1)
        
        # 配置容器的网格权重
        container.grid_columnconfigure(0, weight=2)  # 左侧面板占2份
        container.grid_columnconfigure(1, weight=1)  # 右侧面板占1份
        container.grid_rowconfigure(0, weight=1)
        
        # 创建左右分栏
        left_panel = ttk.Frame(container)
        right_panel = ttk.Frame(container)
        
        # 使用grid布局
        left_panel.grid(row=0, column=0, sticky='nsew', padx=(0, 2))
        right_panel.grid(row=0, column=1, sticky='nsew', padx=(2, 0))
        
        # 配置左右面板的网格权重
        left_panel.grid_rowconfigure(1, weight=1)  # 文件列表区域占据剩余空间
        left_panel.grid_columnconfigure(0, weight=1)
        right_panel.grid_rowconfigure(1, weight=1)  # 模板列表区域占据剩余空间
        right_panel.grid_columnconfigure(0, weight=1)
        
        # === 左侧面板 - 文件列表 ===
        # 文件操作区
        file_frame = ttk.Frame(left_panel)
        file_frame.pack(fill=tk.X, pady=5)
        
        ttk.Button(file_frame, text="添加文件", command=self.add_files).pack(side=tk.LEFT, padx=5)
        ttk.Button(file_frame, text="打开目录", command=self.open_directory).pack(side=tk.LEFT, padx=5)
        ttk.Button(file_frame, text="清空列表", command=self.clear_files).pack(side=tk.LEFT, padx=5)
        
        # 文件列表
        list_frame = ttk.LabelFrame(left_panel, text="待分析文件", padding="5")
        list_frame.pack(fill=tk.BOTH, expand=True, pady=5)
        
        # 添加多选提示
        ttk.Label(
            list_frame, 
            text="提示：按住 Ctrl 键可多选文件，按住 Shift 键可连续选择",
            foreground='gray'
        ).pack(anchor=tk.W, padx=5, pady=(0, 5))
        
        # 创建带滚动条的树形视图
        self.file_list = ttk.Treeview(
            list_frame,
            columns=("文件名", "大小", "状态", "进度"),
            show="headings",
            selectmode="extended"  # 允许多选
        )
        
        # 设置列
        self.file_list.heading("文件名", text="文件名")
        self.file_list.heading("大小", text="大小")
        self.file_list.heading("状态", text="状态")
        self.file_list.heading("进度", text="进度")
        
        self.file_list.column("文件名", width=250)
        self.file_list.column("大小", width=80)
        self.file_list.column("状态", width=200)
        self.file_list.column("进度", width=100)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(list_frame, orient=tk.VERTICAL, command=self.file_list.yview)
        self.file_list.configure(yscrollcommand=scrollbar.set)
        
        self.file_list.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 控制按钮
        button_frame = ttk.Frame(left_panel)
        button_frame.pack(fill=tk.X, pady=5)
        
        ttk.Button(button_frame, text="开始分析", command=self.start_analysis).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="停止", command=self.stop_analysis).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="导出所有结果", command=self.export_all_results).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="删除选中", command=self.remove_selected).pack(side=tk.LEFT, padx=5)  # 添加删除按钮
        
        # === 右侧面板 - 分析设置 ===
        # 角色选择区（移到最上面）
        role_frame = ttk.LabelFrame(right_panel, text="分析师角色", padding="5")
        role_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 主角色选择
        main_role_frame = ttk.Frame(role_frame)
        main_role_frame.pack(fill=tk.X, pady=2)
        ttk.Label(main_role_frame, text="主要角色:").pack(side=tk.LEFT, padx=5)
        main_role_combo = ttk.Combobox(
            main_role_frame,
            textvariable=self.main_role_var,
            values=list(ANALYST_ROLES.keys()),
            state="readonly",
            width=20
        )
        main_role_combo.pack(side=tk.LEFT, padx=5)
        main_role_combo.bind('<<ComboboxSelected>>', self.on_main_role_change)
        
        # 辅助角色选择
        assistant_frame = ttk.LabelFrame(role_frame, text="辅助角色", padding="5")
        assistant_frame.pack(fill=tk.X, pady=2)
        
        self.assistant_roles_listbox = tk.Listbox(
            assistant_frame,
            selectmode=tk.MULTIPLE,
            height=4,
            exportselection=False
        )
        
        # 初始化辅助角色列表
        for role in ANALYST_ROLES.keys():
            if role != self.main_role_var.get():
                self.assistant_roles_listbox.insert(tk.END, role)
                if role in self.selected_assistant_roles:
                    self.assistant_roles_listbox.selection_set(tk.END)
        
        self.assistant_roles_listbox.pack(fill=tk.X, padx=5, pady=5)
        
        # 保存设置区域
        save_frame = ttk.LabelFrame(right_panel, text="保存设置", padding="5")
        save_frame.pack(fill=tk.X, pady=(0, 5))
        
        # 自动保存选项
        auto_save_frame = ttk.Frame(save_frame)
        auto_save_frame.pack(fill=tk.X, pady=2)
        
        self.auto_save_var = tk.BooleanVar(value=self.auto_save)
        ttk.Checkbutton(
            auto_save_frame,
            text="分析完成后自动保存",
            variable=self.auto_save_var
        ).pack(side=tk.LEFT, padx=5)
        
        # 保存格式选择
        format_frame = ttk.Frame(save_frame)
        format_frame.pack(fill=tk.X, pady=2)
        
        ttk.Label(format_frame, text="保存格式:").pack(side=tk.LEFT, padx=5)
        self.save_format_var = tk.StringVar(value=self.save_format)
        format_combo = ttk.Combobox(
            format_frame,
            textvariable=self.save_format_var,
            values=["txt", "markdown"],
            state="readonly",
            width=15
        )
        format_combo.pack(side=tk.LEFT, padx=5)
        format_combo.bind('<<ComboboxSelected>>', lambda e: self.save_settings())
        
        # 保存路径选择
        path_frame = ttk.Frame(save_frame)
        path_frame.pack(fill=tk.X, pady=2)
        
        ttk.Label(path_frame, text="保存位置:").pack(side=tk.LEFT, padx=5)
        self.save_path_var = tk.StringVar(value=self.save_path)
        self.save_path_entry = ttk.Entry(
            path_frame,
            textvariable=self.save_path_var,
            state="readonly"
        )
        self.save_path_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        ttk.Button(
            path_frame,
            text="选择",
            command=self.choose_save_path
        ).pack(side=tk.LEFT, padx=5)
        
        # 模板选择区（移到最下面）
        template_frame = ttk.LabelFrame(right_panel, text="分析模板", padding="5")
        template_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 5))
        
        # 模板列表
        self.template_listbox = tk.Listbox(
            template_frame, 
            selectmode=tk.MULTIPLE,
            exportselection=False
        )
        template_scroll = ttk.Scrollbar(template_frame, orient=tk.VERTICAL, command=self.template_listbox.yview)
        self.template_listbox.configure(yscrollcommand=template_scroll.set)
        
        # 绑定选择事件
        self.template_listbox.bind('<<ListboxSelect>>', self.on_template_select)
        
        self.template_listbox.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(5, 0))
        template_scroll.pack(side=tk.RIGHT, fill=tk.Y, padx=(0, 5))
        
        # 添加模板到列表框
        for template in get_template_names():
            self.template_listbox.insert(tk.END, template)
        
        # 模板操作按钮
        template_button_frame = ttk.Frame(template_frame)
        template_button_frame.pack(fill=tk.X, pady=5)
        ttk.Button(template_button_frame, text="全选", command=self.select_all_templates).pack(side=tk.LEFT, padx=2)
        ttk.Button(template_button_frame, text="取消全选", command=self.deselect_all_templates).pack(side=tk.LEFT, padx=2)
    
    def add_files(self):
        """添加文件到列表"""
        new_files = filedialog.askopenfilenames(
            title="选择要分析的文件",
            filetypes=[
                ("文本文件", "*.txt"),
                ("Word文档", "*.docx"),
                ("PDF文件", "*.pdf"),
                ("所有文件", "*.*")
            ]
        )
        
        if not new_files:
            return
            
        for path in new_files:
            if path not in self.file_paths:
                self.file_paths.append(path)
                size = self.get_file_size(path)
                self.file_list.insert("", tk.END, values=(
                    os.path.basename(path),
                    size,
                    "等待中",
                    "0%"
                ))
    
    def clear_files(self):
        """清空文件列表"""
        if messagebox.askyesno("确认", "确定要清空文件列表吗？"):
            self.file_list.delete(*self.file_list.get_children())
            self.file_paths = []
            self.results = {}
    
    def get_file_size(self, file_path):
        """获取文件大小的人类可读格式"""
        size = os.path.getsize(file_path)
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"
    
    def read_file_content(self, file_path):
        """根据文件类型读取内容"""
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        elif ext == '.docx':
            doc = Document(file_path)
            return '\n'.join(p.text for p in doc.paragraphs)
        elif ext == '.pdf':
            reader = PdfReader(file_path)
            return '\n'.join(page.extract_text() for page in reader.pages)
        else:
            raise ValueError(f"不支持的文件类型: {ext}")
    
    async def analyze_file(self, file_path, item_id):
        """分析单个文件"""
        try:
            # 更新状态
            self.file_list.set(item_id, "状态", "正在读取文件...")
            content = self.read_file_content(file_path)
            
            self.file_list.set(item_id, "状态", "正在分析...")
            # 使用当前选择的角色和模板
            prompt = combine_templates(
                self.get_selected_templates(),
                content,
                self.main_role_var.get(),
                self.get_selected_assistant_roles()
            )
            
            # 分析文本
            result = ""
            total_chunks = 0
            async for output_type, chunk in self.api.analyze_text_stream(content, prompt):
                if not self.is_running:
                    raise InterruptedError("分析已停止")
                    
                result += chunk
                total_chunks += 1
                
                # 更新进度（每10个chunk更新一次）
                if total_chunks % 10 == 0:
                    self.file_list.set(item_id, "进度", f"已处理 {total_chunks} 块")
            
            # 保存结果
            self.results[file_path] = result
            self.file_list.set(item_id, "状态", "完成")
            self.file_list.set(item_id, "进度", "100%")
            
            # 自动保存
            if self.auto_save_var.get():
                await self.save_result(file_path, result)
            
        except InterruptedError as e:
            self.file_list.set(item_id, "状态", "已停止")
            self.file_list.set(item_id, "进度", "-")
        except Exception as e:
            self.file_list.set(item_id, "状态", f"错误: {str(e)}")
            self.file_list.set(item_id, "进度", "失败")
    
    def start_analysis(self):
        """开始批量分析"""
        if self.is_running:
            return
        
        # 获取选中的文件
        selected_items = self.file_list.selection()
        if not selected_items:
            if not messagebox.askyesno("确认", "没有选中任何文件，是否分析所有文件？"):
                return
        
        if not self.get_selected_templates():
            messagebox.showerror("错误", "请至少选择一个分析模板")
            return
        
        self.is_running = True
        threading.Thread(target=self._run_analysis, args=(selected_items,), daemon=True).start()
    
    def stop_analysis(self):
        """停止分析"""
        if self.is_running:
            if messagebox.askyesno("确认", "确定要停止分析吗？"):
                self.is_running = False
    
    def _run_analysis(self, selected_items=None):
        """在新线程中运行分析"""
        async def run():
            items = selected_items if selected_items else self.file_list.get_children()
            for item in items:
                if not self.is_running:
                    break
                
                file_name = self.file_list.item(item)['values'][0]
                file_path = next(
                    path for path in self.file_paths 
                    if os.path.basename(path) == file_name
                )
                
                await self.analyze_file(file_path, item)
        
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(run())
        finally:
            loop.close()
            self.is_running = False
    
    def export_all_results(self):
        """导出所有分析结果"""
        if not self.results:
            messagebox.showwarning("提示", "没有可导出的结果")
            return
        
        export_dir = filedialog.askdirectory(title="选择导出目录")
        if not export_dir:
            return
        
        try:
            for file_path, result in self.results.items():
                base_name = os.path.splitext(os.path.basename(file_path))[0]
                result_path = os.path.join(export_dir, f"{base_name}_分析结果.txt")
                
                with open(result_path, 'w', encoding='utf-8') as f:
                    f.write(result)
            
            messagebox.showinfo("成功", "所有结果已导出")
        except Exception as e:
            messagebox.showerror("错误", f"导出结果时出错: {str(e)}")
    
    def select_all_templates(self):
        """全选所有模板"""
        self.template_listbox.select_set(0, tk.END)
    
    def deselect_all_templates(self):
        """取消全选"""
        self.template_listbox.selection_clear(0, tk.END)

    def get_selected_templates(self):
        """获取选中的模板名称列表"""
        selections = self.template_listbox.curselection()
        return [self.template_listbox.get(i) for i in selections]

    def get_selected_assistant_roles(self):
        """获取选中的辅助角色列表"""
        selections = self.assistant_roles_listbox.curselection()
        return [self.assistant_roles_listbox.get(i) for i in selections]

    def choose_save_path(self):
        """选择保存路径"""
        path = filedialog.askdirectory(title="选择保存目录")
        if path:
            self.save_path_var.set(path)
            self.save_settings()

    async def save_result(self, source_file: str, content: str):
        """保存分析结果"""
        try:
            # 确定保存路径
            save_dir = self.save_path_var.get()
            if save_dir == "与源文件相同目录":
                save_dir = os.path.dirname(source_file)
            
            # 生成保存文件名
            base_name = os.path.splitext(os.path.basename(source_file))[0]
            ext = ".md" if self.save_format_var.get() == "markdown" else ".txt"
            save_path = os.path.join(save_dir, f"{base_name}_分析结果{ext}")
            
            # 根据格式保存
            if self.save_format_var.get() == "markdown":
                content = f"""# {base_name} 分析报告

## 基本信息
- 分析时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
- 源文件：{os.path.basename(source_file)}
- 分析模板：{", ".join(self.get_selected_templates())}
- 主分析师：{self.main_role_var.get()}
- 辅助分析师：{", ".join(self.get_selected_assistant_roles())}

## 分析结果
{content}

---
*本报告由 AI 分析系统生成，仅供参考。*
"""
            else:
                content = f"""分析报告 - {base_name}
{'='*50}

基本信息：
- 分析时间：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
- 源文件：{os.path.basename(source_file)}
- 分析模板：{", ".join(self.get_selected_templates())}
- 主分析师：{self.main_role_var.get()}
- 辅助分析师：{", ".join(self.get_selected_assistant_roles())}

分析结果：
{'-'*50}
{content}

{'='*50}
注：本报告由 AI 分析系统生成，仅供参考。
"""
            
            # 保存文件
            with open(save_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            # 更新状态
            self.file_list.set(
                self.file_list.get_children()[-1],
                "状态",
                f"完成 (已保存到: {save_path})"
            )
            
        except Exception as e:
            raise Exception(f"保存结果时出错: {str(e)}")

    def on_main_role_change(self, event=None):
        """当主角色改变时更新辅助角色列表"""
        # 保存当前选中的辅助角色
        self.selected_assistant_roles = self.get_selected_assistant_roles()
        
        # 更新辅助角色列表
        self.update_assistant_roles()

    def update_assistant_roles(self):
        """更新辅助角色列表"""
        # 保存当前选中的角色
        selected_roles = self.get_selected_assistant_roles()
        
        # 清空并重新填充列表
        self.assistant_roles_listbox.delete(0, tk.END)
        
        # 添加角色并恢复选中状态
        for role in ANALYST_ROLES.keys():
            if role != self.main_role_var.get():
                self.assistant_roles_listbox.insert(tk.END, role)
                # 如果这个角色之前被选中，就选中它
                if role in selected_roles:
                    # 获取当前角色的索引
                    idx = self.assistant_roles_listbox.size() - 1
                    self.assistant_roles_listbox.selection_set(idx)

    def on_template_select(self, event=None):
        """当选择模板时的处理"""
        # 不需要做任何处理，因为我们使用了 exportselection=False
        pass

    def open_directory(self):
        """打开目录并添加所有txt文件"""
        directory = filedialog.askdirectory(title="选择包含文本文件的目录")
        if not directory:
            return
        
        # 遍历目录及其子目录
        txt_files = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                if file.lower().endswith('.txt'):
                    full_path = os.path.join(root, file)
                    txt_files.append(full_path)
        
        if not txt_files:
            messagebox.showinfo("提示", "所选目录下没有找到txt文件")
            return
        
        # 添加找到的文件到列表
        added_count = 0
        for path in txt_files:
            if path not in self.file_paths:
                self.file_paths.append(path)
                size = self.get_file_size(path)
                self.file_list.insert("", tk.END, values=(
                    os.path.basename(path),
                    size,
                    "等待中",
                    "0%"
                ))
                added_count += 1
        
        if added_count > 0:
            messagebox.showinfo("成功", f"已添加 {added_count} 个txt文件")

    def load_settings(self):
        """加载保存的设置"""
        try:
            if os.path.exists('batch_analyzer_settings.json'):
                with open('batch_analyzer_settings.json', 'r', encoding='utf-8') as f:
                    settings = json.load(f)
                    self.save_format = settings.get('save_format', 'txt')
                    self.save_path = settings.get('save_path', '与源文件相同目录')
                    self.auto_save = settings.get('auto_save', True)
            else:
                self.save_format = 'txt'
                self.save_path = '与源文件相同目录'
                self.auto_save = True
        except:
            self.save_format = 'txt'
            self.save_path = '与源文件相同目录'
            self.auto_save = True

    def save_settings(self):
        """保存设置"""
        settings = {
            'save_format': self.save_format_var.get(),
            'save_path': self.save_path_var.get(),
            'auto_save': self.auto_save_var.get()
        }
        try:
            with open('batch_analyzer_settings.json', 'w', encoding='utf-8') as f:
                json.dump(settings, f, ensure_ascii=False, indent=2)
        except:
            pass

    def remove_selected(self):
        """删除选中的文件"""
        selected_items = self.file_list.selection()
        if not selected_items:
            return
        
        if messagebox.askyesno("确认", "确定要删除选中的文件吗？"):
            for item in selected_items:
                # 获取文件名
                file_name = self.file_list.item(item)['values'][0]
                # 从文件路径列表中移除
                self.file_paths = [path for path in self.file_paths 
                                 if os.path.basename(path) != file_name]
                # 从显示列表中移除
                self.file_list.delete(item)
                # 从结果字典中移除
                for path in list(self.results.keys()):
                    if os.path.basename(path) == file_name:
                        del self.results[path] 